open Format
open Ast
open Datatypes
open Locations

let fmt = std_formatter;;

(* keep track of the startnode number *)
let startNum = ref 0;;

let getStartnodeNum () = startNum := (!startNum + 1); !startNum;;

(* sets out channel, standard is standard out. If it's a file, this sets the file name and the file ending according to the format given *)
let setFormatterOutChannel fileEnding = function
	Some filename -> let out = open_out_gen [ Open_text; Open_creat; Open_excl; Open_wronly ] 0o644 (filename ^ "." ^ fileEnding) in set_formatter_out_channel out
	| None -> ()
	;;
	

let printStr fmt s= fprintf fmt "%s" s;;

let postLabelStr labels _ = 
  let Labels(_, post) = labels in
    match post with
      Post(l) -> " ]:l" ^ string_of_int l ^ "_"
      | _ -> raise (Label_missmatch ("A labels pair must be on the form (Prelabel, Postlabel)."))
  ;;
  
let preLabelStr labels _ = 
  let Labels(pre, _) = labels in
    match pre with
      Pre(l) -> "l" ^ string_of_int l ^ ":[ "
      | _ -> raise (Label_missmatch ("A labels pair must be on the form (Prelabel, Postlabel)."))
  ;;
  
let emptyStr _ _ = "";;

(* takes a map from labels to locations and *)
let postLocationStr labelMap (Labels(_,post)) = function
	("comp" | "declcomp" | "proc") -> ""
	| "last" -> let Location(l) = getLocation post labelMap in (" :L" ^ (string_of_int l))
	| "declblock" -> let Location(l) = getLocation post labelMap in (" :L" ^ (string_of_int l))
	| _ -> " ]"
	;;
	
let preLocationStr labelMap (Labels(pre, _)) = function
	("comp" | "declcomp" | "proc") -> ""
	| _ -> let Location(l) = getLocation pre labelMap in ("L" ^ (string_of_int l) ^ ":" ^ "[ ")
	;;

let rec printList printer sep fmt = function
  |[h] -> fprintf fmt "%a" printer h
  |h::t ->
      fprintf fmt "%a%s@ %a"
      printer h sep (printList printer sep) t
  |[] -> ();;
  
let rec printValue fmt = function
  True -> fprintf fmt "%s" "tt"
  | False -> fprintf fmt "%s" "ff"
  ;;
  
let rec printExp fmt = function
	Value(True) -> fprintf fmt "%s" "tt"
	| Value(False) ->  fprintf fmt "%s" "ff"
	| Variable(id) ->  fprintf fmt "%s" id
	| Not(e) ->  fprintf fmt "%s%a" "!" printExp e
	| Equal(e1, e2) -> fprintf fmt "%a%s%a" printExp e1 " = " printExp e2
	| And(e1, e2) -> fprintf fmt "%a%s%a" printExp e1 " & " printExp e2
	| Or(e1, e2) -> fprintf fmt "%a%s%a" printExp e1 " | " printExp e2
	| Group(e) -> fprintf fmt "%s%a%s" "(" printExp e ")"
	;;
  
let printDecider fmt = function
  Undet -> fprintf fmt "%s" "*"
  | Expression exp -> printExp fmt exp
  ;;

let rec printStm prep postp fmt = function
	Skip labels -> fprintf fmt "%s%s%s" (prep labels "skip") "skip" ((postp labels "skip") ^ ";")
	| Comp(labels, s1, s2) -> fprintf fmt "%s%a@\n%a%s" (prep labels "comp") (printStm prep postp) s1 (printStm prep postp)  s2 (postp labels "comp")
	| Callvoid(labels, id1, eList) -> fprintf fmt "%s%s%a%s%s" (prep labels "callvoid") (id1 ^ " (") (printList printExp ",") eList ")" ((postp labels "callvoid") ^ ";")
	| Callbool(labels, id1, id2, eList) -> fprintf fmt "%s%s%a%s%s" (prep labels "callbool") (id1 ^ " := " ^ id2 ^ " (") (printList printExp ",") eList ")" ((postp labels "callbool") ^ ";")
	| Assignment(labels, idL, eList) -> fprintf fmt "%s%a%s%a%s" (prep labels "ass") (printList printStr ",") idL " := " (printList printExp ",") eList ((postp labels "ass") ^ ";")
	| While(labels, d, s) -> fprintf fmt "@[<hov3>%s%s%a%s@\n%a@]@\n%s%s" (prep labels "while") "while (" printDecider d ") do" (printStm prep postp) s "od" ((postp labels "while") ^ ";")
	| Returnvoid(labels) -> fprintf fmt "%s%s%s" (prep labels "returnvoid") "return" ((postp labels "returnvoid") ^ ";")
	| Returnbool(labels, e) -> fprintf fmt "%s%s%a%s" (prep labels "returnbool") "return " printExp e ((postp labels "returnbool") ^ ";")
	| Ifelse(labels, d, s1, s2) -> fprintf fmt "@[<hov3>%s%s%a%s@\n%a@]@\n%s@[<hov-1>@\n%a@]@\n%s%s" (prep labels "if") "if (" printDecider d ") then" (printStm prep postp) s1 "else" (printStm prep postp) s2 "fi" ((postp labels "if") ^ ";")
	;;

let rec printDecl prep postp fmt = function
	Declaration(labels, id) -> fprintf fmt "%s%s%s%s%s" (prep labels "decl") "decl " id (postp labels "decl") ";"
	| Declcomp(labels, d1, d2) -> fprintf fmt "%s%a@\n%a%s" (prep labels "declcomp") (printDecl prep postp) d1 (printDecl prep postp) d2 (postp labels "declcomp")
	| Empty -> fprintf fmt "%s" ""
	;;

let printProctype fmt = function
		Bool -> fprintf fmt "%s" "bool"
		| Void -> fprintf fmt "%s" "void"
		;;
		
(* Prints a single procedure *)	
let printProc prep postp fmt proc = let Procedure(labels, proctype, id, fp, d, s) = proc in
                                    if (d <> Empty) then begin
                                      fprintf fmt "@[<hov3>%s%a%s%s%a%s%s@\n%a@\n%a@]@\n%s"
                                      (prep labels "proc")
                                      printProctype proctype
                                      (" " ^ id)
                                      " (" 
                                      (printList printStr ",") fp
                                      ") "
                                      "begin"
                                      (printDecl prep postp) d
                                      (printStm prep postp) s
                                      "end"
                                    end else begin
                                      fprintf fmt "@[<hov3>%s%a%s%s%a%s%s@\n%a@]@\n%s"
                                      (prep labels "proc")
                                      printProctype proctype
                                      (" " ^ id)
                                      " (" 
                                      (printList printStr ",") fp
                                      ") "
                                      "begin"
                                      (printStm prep postp) s
                                      "end"
                                    end
                                    ;;
                                    
let	rec printProcs prep postp fmt = function
	| [p] -> let Procedure(labels, _, _, _, _, _) = p in fprintf fmt "%a%s" (printProc prep postp) p (postp labels "declblock")
	| p::ps -> let Procedure(labels, _, _, _, _, _) = p in fprintf fmt "%a%s@\n@\n%a" (printProc prep postp) p (postp labels "declblock") (printProcs prep postp) ps
	| [] -> ()
	;;

	
let printLabel fmt = function
	Post l -> fprintf fmt "%s%s" "_l" (string_of_int l)
	| Pre l -> fprintf fmt "%s%s" "l" (string_of_int l)
	;;
	
let printCN fmt cn =
	let ControlNode(num) = cn in
		fprintf fmt "%s%s" "v" (string_of_int num)
	;;
	
let printLocation fmt (Location(l)) = fprintf fmt "%s%d" "L" l;;

let printLocationForLabel fmt lablocPair = let (lab, loc) = lablocPair in fprintf fmt "%a%s%a@\n" printLabel lab " -> " printLocation loc;;

let compareLabels lab1 lab2 = match lab1 with
	Pre(i1) -> (match lab2 with
		Pre(i2) -> compare i1 i2
		| Post(i2) -> compare i1 i2)
	| Post(i1) -> (match lab2 with
		Pre(i2) -> compare i1 i2
		| Post(i2) -> compare i1 i2)
	;;
	
let compareLabelMap (lab1, _) (lab2, _) = compareLabels lab1 lab2;;

let printLabelMap fmt labMap = List.iter (printLocationForLabel fmt) (List.sort compareLabelMap labMap);;

let printLabelsForLocation fmt loclabPair = let (loc, labs) = loclabPair in fprintf fmt "%a%s%a%s@\n" printLocation loc " -> { " (printList printLabel ",") labs " }"

let printLocationMap fmt locMap = List.iter (printLabelsForLocation fmt) (List.sort compareLocMap locMap);;

	
let printWithLabels decl procs = match decl with
	Empty -> fprintf fmt "%a"(printProcs preLabelStr postLabelStr) procs
	| _ -> fprintf fmt "%a@\n%a" (printDecl preLabelStr postLabelStr) decl (printProcs preLabelStr postLabelStr) procs
	;;

let printWithLocations decl procs = match decl with
	Empty -> let (labMap, locMap) = Locations.createLocations (Bp(decl, procs)) in fprintf fmt "%a" (printProcs (preLocationStr labMap) (postLocationStr labMap)) procs (*printLabelMap fmt labMap; printLocationMap fmt locMap*)
	| Declaration(labels, string) -> 
		let (labMap, locMap) = Locations.createLocations (Bp(decl, procs)) in 
			fprintf fmt "%a%s@\n@\n%a" (printDecl (preLocationStr labMap) (postLocationStr labMap)) decl 
									(postLocationStr labMap labels "declblock")
									(printProcs (preLocationStr labMap) (postLocationStr labMap)) procs; printLabelMap fmt labMap; printLocationMap fmt locMap
	| Declcomp(labels,_,_) ->
		let (labMap, locMap) = Locations.createLocations (Bp(decl, procs)) in 
			fprintf fmt "%a%s@\n@\n%a" (printDecl (preLocationStr labMap) (postLocationStr labMap)) decl 
									(postLocationStr labMap labels "declblock")
									(printProcs (preLocationStr labMap) (postLocationStr labMap)) procs; printLabelMap fmt labMap; printLocationMap fmt locMap
	;;
	(*
let printWithLocations decl procs = match decl with
	Empty -> let (labMap, locMap) = Locations.createLocations (Bp(decl, procs)) in printLabelMap fmt labMap; printLocationMap fmt locMap
	| _ -> let (labMap, locMap) = Locations.createLocations (Bp(decl, procs)) in printLabelMap fmt labMap; printLocationMap fmt locMap
	;;
*)	

let printCode decl procs = match decl with
	Empty -> fprintf fmt "%a"(printProcs emptyStr emptyStr) procs
	| _ -> fprintf fmt "%a@\n%a" (printDecl emptyStr emptyStr) decl (printProcs emptyStr emptyStr) procs
	;;

	
let rec printCMap fmt = function (* match CMaps *)
	fst :: rest -> 
		let (loc, code) = fst in
			(match code with
				CodeStm(s) -> fprintf fmt "%a%s%a@\n" printLocation loc " -> " (printStm emptyStr emptyStr) s; printCMap fmt rest
				| CodeDecl(d) -> fprintf fmt "%a%s%a@\n" printLocation loc " -> " (printDecl emptyStr emptyStr) d; printCMap fmt rest
        | Nothing -> ()
      )
	| [] -> ()
	;;

let rec printState fmt = function (* match state: (string, expression) *)
  [] -> ()	
	| fst :: rest ->
		let (var, value) = fst in
			fprintf fmt "%s%s%a%s"  var " -> " printValue value ", ";
		printState fmt rest
	;;


let rec printWCNTriples fmt = function (* match wraped control node triples (cn, label, state) *)
	(fst :: rest) ->
    (match fst with
    (Visited(cn, label, state) | Unvisited(cn, label, state)) ->
      fprintf fmt "%s%a%s%a%s%a%s@\n" "( " printCN cn ", " printLabel label " -> [ " printState state " ] )";
		printWCNTriples fmt rest)
	| [] -> ()
	;;
  
let rec printCNTriples fmt = function (* match control node triples (cn, label, state) *)
	(cn, label, state) :: rest ->
      (fprintf fmt "%s%a%s%a%s%a%s@\n" "( " printCN cn ", " printLabel label " -> [ " printState state " ] )";
      printCNTriples fmt rest)
	| [] -> ()
	;;
	
let rec printTransitions fmt = function (* match transitions *)
	fst :: rest -> 
		let (cn1, a, cn2) = fst in
			fprintf fmt "@[<h>%s@ %a@ %a@ %s@]@." "edge" printCN cn1 printCN cn2 a;
			printTransitions fmt rest
	| [] -> ()
	;;	

  (* prints the atomic propositions for a node, identical to printList with sep = "" except that this function replaces strins x equal to pName with meth(x) to mark that this atomic proposition is a method name for the current method*)
let rec printValuation pName fmt = function
  |[h] -> 
    if (pName = h) then
      fprintf fmt "%s%a%s" "meth(" printStr h ")"
    else
      fprintf fmt "%a" printStr h
  |h::t ->
    if (pName = h) then
      fprintf fmt "%s%a%s@ %a" "meth(" printStr h ")" (printValuation pName) t
    else
      fprintf fmt "%a@ %a" printStr h (printValuation pName) t;
  |[] -> ();;
  
(* valuations = [(cn, [string])] *)
let rec printNodes pName entryCount fmt = function (* match valuations *)
	fst :: rest -> 
		let (cn, pList) = fst in
      if (entryCount <= 0) then begin
        fprintf fmt "@[<h>%s@ %a@ %a@]@." "node" printCN cn (printValuation pName) pList;
        printNodes pName 0 fmt rest
      end else begin
        fprintf fmt "@[<h>%s@ %a@ %a@ %s@]@." "node" printCN cn (printValuation pName) pList "entry";
        printNodes pName (entryCount-1) fmt rest
      end
	| [] -> ()
	;;	
 
let printFlowGraph decl procs =
	let bp = Bp(decl, procs) in
	let (labMap, locMap) = Locations.createLocations bp in
	let glob = Translation.getGlobals bp in
	let rec printNodeIter allTransitions procs =
		match procs with
		 proc :: rest ->
			let Procedure(_, _, pName, _, _, _) = proc in
      let realGlob = Translation.getGlobalsInProc glob proc in (* removes the local vairables and formal parameters from the global state *)
			let locToCMap = Translation.createLocToCodeMap labMap proc in
			let entryNodes = Translation.createEntryNodes labMap proc realGlob in
			let (trans, valuations) = Translation.createTVPair pName realGlob locToCMap labMap entryNodes in
				fprintf fmt "%a"
          (printNodes pName (List.length entryNodes))((List.sort (fun (ControlNode(i1), _) (ControlNode(i2), _) -> compare i1 i2)) valuations) (* sorting, remove if this is not importent. *)
          ;
				printNodeIter (trans @ allTransitions) rest
		| [] -> allTransitions
	in
		let allTransitions = printNodeIter [] procs in
      printTransitions fmt (List.sort (fun (ControlNode(i1), _, _) (ControlNode(i2), _, _) -> compare i1 i2) allTransitions)
	;;
	
(* dot printer *)

(* prints num number of start nodes *)
let rec printStartNodes fmt num = 
	if (num > 0) then begin
		fprintf fmt "%s%s%s@\n" "start" (string_of_int (getStartnodeNum ())) " [shape=none, label=\" \"]";
		printStartNodes fmt (num - 1)
	end else
		()
	;;
	
let rec printDotNodes fmt = function (* match valuations *)
	fst :: rest -> 
		let (cn, pList) = fst in
			fprintf fmt "%a%s%a%s%a%s%a%s" 
				printCN cn
				" [label=<<TABLE BORDER=\"0\" CELLBORDER=\"0\" CELLSPACING=\"10\"><TR><TD BORDER=\"1\" PORT=\""
				printCN cn
				"\">"
				printCN cn
				"</TD><TD BORDER=\"0\">"
				(printList printStr ",") pList
				"</TD></TR></TABLE>>]";
			printDotNodes fmt rest
	| [] -> ()
	;;	
	
let rec printDotTransitions fmt = function (* match transitions*)
	fst :: rest -> 
		let (cn, label, nextCN) = fst in
			(fprintf fmt "%a%s%a%s%a%s%a%s%s%s@\n" 
				printCN cn
				":"
				printCN cn
				" -> "
				printCN nextCN
				":"
				printCN nextCN
				"[label = \""
				label
				"\"];";
			printDotTransitions fmt rest)
	| [] -> ()
	;;	

let rec printDotStartTransitions num fmt = function (* match entryNodes *)
	fst :: rest -> 
    (match fst with
    (Visited(cn, _, _) | Unvisited(cn, _, _)) ->
			(fprintf fmt "%s%s%s%a%s%a%s@\n"
				"start"
				(string_of_int (!startNum - num + 1))
				" -> "
				printCN cn
				":"
				printCN cn
				";";
			printDotStartTransitions (num-1) fmt rest))
	| [] -> ()
	;;

(* print a graph represented in Dot *)
let printDot decl procs =
	let bp = Bp(decl, procs) in
	let (labMap, locMap) = Locations.createLocations bp in
	let glob = Translation.getGlobals bp in
	let rec iterateProcs num fmt procs =
		match procs with
		 proc :: rest ->
			let Procedure(_, _, pName, _, _, _) = proc in
      let realGlob = Translation.getGlobalsInProc glob proc in (* removes the local vairables and formal parameters from the global state *)
			let locToCMap = Translation.createLocToCodeMap labMap proc in
			let entryNodes = Translation.createEntryNodes labMap proc realGlob in
			let (trans, valuations) = Translation.createTVPair pName realGlob locToCMap labMap entryNodes in
				fprintf fmt "%s%s%s@\n%s@\n%s@\n%a%a%a%a@\n%s@\n@\n"
					"subgraph cluster"
					(string_of_int num)
					" {"
					("label=\"Method: " ^ pName ^"\";") 
					"node [shape=plaintext]"
					printStartNodes (List.length entryNodes)
					printDotNodes valuations
					printDotTransitions trans
					(printDotStartTransitions (List.length entryNodes)) entryNodes
					"}";
				iterateProcs (num + 1) fmt rest
		| [] -> ()
	in
		fprintf fmt "%s@\n%a@\n%s"
			"digraph structs {"
			(iterateProcs 1) procs
			"}"
	;;
  
  let printPreprocessed decl procs =
    let bp = Bp(decl, procs) in
      let glob = Translation.getGlobals bp in
      let newProcs = Preprocess.prep procs glob in
      begin
        printCode decl newProcs
      end
    ;;